import logging
import os
import json
import shutil
from typing import List, Dict, Any, Optional, Tuple, Set
from source.data_processor.utils import ExperimentData
from .experiment import comparison
from .experiment.comparison.models import get_models
from .experiment.comparison.utils import ModelParams
from .model.graph import Graph
from .model.case import Case
from .ExperimentCaseConverter import ExperimentCaseConverter
from ..graph.casualGraph import CausalGraph


# 静态图工厂类，用于封装单个图并提供 CircaRunner 所需的接口
class StaticGraphFactory:
    """
    静态图工厂类，用于封装内存图实例。
    """

    def __init__(self, graph: Graph):
        self.graph = graph

    def get_graph(self) -> Graph:
        """返回封装的图实例"""
        return self.graph


# 从 experiment/__main__.py 引入的常量
_GRAPH_SIZES: List[Tuple[int, int]] = [
    (50, 100),
    (100, 500),
    (500, 5000),
]
_NUM_GRAPHS = 10
_DEFAULT_SEED = 519


class CircaRunner:
    """
    CircaRunner 类用于运行 CIRCA 实验。
    使用 ExperimentData 和 CausalGraph 直接运行，基于 Case 对象而非 SimDataset。
    """

    def __init__(
        self,
        output_dir: str = "./output",
        report_dir: str = "./report",
        model_params_path: str = "D:\project\k8s\iscas-rca\iscas-rca-Jayx\code\localize\circa\params-sim-tune.json",
        seed: int = _DEFAULT_SEED,
        cuda: bool = True,
        max_workers: int = 1,
        comparison_delay: Optional[int] = None,
        # 实验数据和图信息
        experiment_data: Optional[ExperimentData] = None,
        causal_graph: Optional[CausalGraph] = None,
        # 真实根因
        ground_truth_causes: Optional[Set[Tuple[str, str]]] = None,
        # 覆盖 SimCase 中的 case_data_params
        case_data_params_override: Optional[Dict[str, Any]] = None,
    ):
        self.output_dir = output_dir
        self.report_dir = report_dir
        self.seed = seed
        self.cuda = cuda
        self.max_workers = max_workers
        self.comparison_delay = comparison_delay

        # 创建日志记录器
        self.logger = logging.getLogger(self.__class__.__name__)

        # 删除 output_dir 下的所有内容（如果存在）
        if os.path.exists(self.output_dir):
            shutil.rmtree(self.output_dir)
        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(self.report_dir, exist_ok=True)

        try:
            with open(model_params_path, "r", encoding="utf-8") as f:
                params_json = json.load(f)
            self.model_params = ModelParams(params_json)
        except Exception as e:
            self.logger.error(f"从 {model_params_path} 加载 ModelParams 失败: {e}")
            self.model_params = ModelParams()

        self.experiment_data = experiment_data
        self.causal_graph = causal_graph
        self.ground_truth_causes = ground_truth_causes
        self.case_data_params_override = case_data_params_override

    def run_experiment(self):
        """
        运行 CIRCA 实验。直接使用 Case 对象。
        """
        if (
            not self.experiment_data
            or not self.causal_graph
            or not self.ground_truth_causes
        ):
            self.logger.error("实验所需的一个或多个关键对象未设置。")
            return

        converter = ExperimentCaseConverter(
            experiment_data=self.experiment_data,
            causal_graph=self.causal_graph,
            ground_truth_causes=self.ground_truth_causes,
        )

        try:
            # 使用create_case方法获取Case对象
            case: Case = converter.create_case()
            # 获取内存图用于创建图工厂
            memory_graph = case.graph
            graph_factory = StaticGraphFactory(graph=memory_graph)
        except Exception as e:
            self.logger.error(f"构建Case失败: {e}")
            return

        # 创建输出和报告路径
        anomaly_id_str = (
            str(self.experiment_data.anomaly_id).replace("/", "_").replace(":", "_")
        )
        current_output_dir = os.path.join(self.output_dir)
        os.makedirs(current_output_dir, exist_ok=True)
        current_report_filename = os.path.join(
            self.report_dir, f"report-case_mode-{anomaly_id_str}.csv"
        )

        self._process_case(
            case,
            graph_factory,
            current_output_dir,
            current_report_filename,
            identifier=f"exp-{anomaly_id_str}",
        )

    def _process_case(
        self,
        case: Case,
        graph_factory: StaticGraphFactory,
        output_dir: str,
        report_filename: str,
        identifier: str,
    ):
        """
        处理单个 Case 对象。
        """
        models, graph_factories = get_models(
            graph_factories={"GT": graph_factory},
            params=self.model_params,
            seed=self.seed,
            cuda=self.cuda,
            max_workers=self.max_workers,
        )

        comparison_args = {
            "models": models,
            "cases": [case],  # 使用单个 Case 对象的列表
            "output_dir": output_dir,
            "report_filename": report_filename,
            "max_workers": 1,
        }

        if self.comparison_delay is not None:
            comparison_args["delay"] = self.comparison_delay

        try:
            comparison.run(**comparison_args)
        except Exception as e:
            self.logger.error(f"比较运行期间出错: {e}")
            raise e
